import torch
from torch import nn

from utilities import cos_sim

class GCN(nn.Module):
    def __init__(self, input_dim:int, output_dim:int) -> None:
        '''
        input_dim:特征的维度
        output_dim:输出特征的维度
        '''
        super().__init__()
        self.W = nn.Linear(input_dim, output_dim)
        self.B = nn.Linear(input_dim, output_dim)
        self.activate = nn.ReLU()
    def forward(self, feature:torch.Tensor, graph:torch.Tensor) -> torch.Tensor:
        feature = torch.matmul(graph, feature)
        output = self.W(feature) + self.B(feature)
        return self.activate(output), graph

class GCNSelfLoss(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.loss = nn.L1Loss()
    def forward(self,input:torch.Tensor, adjacency:torch.Tensor):
        '''
        input为结点特征
        adjacency为邻接矩阵
        '''
        input = cos_sim(input, input)
        input = input * adjacency
        input = torch.sum(input, -1)
        target = torch.sum(adjacency, -1)
        return self.loss(input, target)

class GCNs(nn.Module):
    def __init__(self, input_dim:int, output_dim:int, deepth:int=2) -> None:
        '''
        input_dim:特征的维度
        output_dim:输出特征的维度
        '''
        super().__init__()
        self.linear = nn.Linear(input_dim, output_dim)
        self.nets = nn.ModuleList()
        for i in range(deepth):
            self.nets.append(GCN(output_dim, output_dim))
        self.activate = nn.Sigmoid()
        
    @staticmethod
    def get_graph(adjacency:torch.Tensor):
        '''
        adjacency:torch.Tensor 为邻接矩阵
        '''
        # 获取度矩阵
        D = torch.sum(adjacency, dim=-1)
        D = D * torch.eye(adjacency.shape[0])
        # D_tmp = torch.linalg.inv(D)**(1/2)
        # 因为有的结点(碱基子串)可能不存在，
        # 于是度为0,对角线存在0,所以不能直接求逆
        # D**(-1/2)会出现inf，暂时没找到替换inf的高效方法
        for i in range(D.shape[0]):
            if D[i][i] != 0:
                D[i][i] = D[i][i]**(-1/2)
        # 计算权重图
        graph = D @ adjacency @ D
        return graph

    def forward(self, feature:torch.Tensor, graph:torch.Tensor) -> torch.Tensor:
        '''
        graph应该包含结点关系信息
        '''
        feature = self.linear(feature)
        for net in self.nets:
            feature, graph = net(feature, graph)
        return self.activate(feature)
